#!/bin/python

from scipy.integrate import ode, cumtrapz
import pickle
import seaborn as sb
import numpy as np
from itertools import product
from figure_2b import figure_2b

def fieldeqns(y, t, epsilont, epsilon, chi, kappa1, kappa2, Delta_a, Delta_b, a, b, c):
    dy = np.array([[0.],[0.]], dtype='complex128')
    
    chi = np.dot(- chi, ((- 1) ** a + (- 1) ** b + (- 1) ** c))
    eps = np.interp(t, epsilont, epsilon)
    
    # Interpolate the data set (gt,g) at time t
    dy[0] = np.dot(- np.sqrt(kappa1), eps) - np.dot(np.dot(1j, (Delta_a + chi)), y[0]) - np.dot(np.dot(0.5, kappa1), y[0]) - np.dot(np.dot(0.5, np.sqrt(np.dot(kappa1, kappa2))), y[1])
    dy[1] = np.dot(- np.sqrt(kappa2), eps) - np.dot(np.dot(1j, (Delta_b + chi)), y[1]) - np.dot(np.dot(0.5, kappa2), y[1]) - np.dot(np.dot(0.5, np.sqrt(np.dot(kappa1, kappa2))), y[0])
    return dy

if __name__ == "__main__":
    
    with open('../python/trajectory_input_file', 'r') as file:
        input_dict = pickle.load(file)

    locals().update(input_dict)

    alpha = np.zeros(shape=(N, 2, 2, 2), dtype='complex128')
    beta = np.zeros(shape=(N, 2, 2, 2), dtype='complex128')

    dt = tvec[1]-tvec[0]
    for a, b, c in product(range(2), repeat=3):
        sol = np.zeros((N, 2), dtype='complex128')
        
        solver = ode(lambda t, y: fieldeqns(y, t, tvec,
            epsilon, chi, kappa1, kappa2, Delta_a, Delta_b, a, b,
            c)).set_integrator('zvode')
        
        solver.set_initial_value([0.+0j,0.+0j],0.0)
        
        sol_pos = 0
        while solver.successful() and solver.t < tvec[-1]:
            sol[sol_pos, :] = solver.y
            solver.integrate(solver.t + dt)
            sol_pos = sol_pos + 1
        
        alpha[:, a, b, c] = sol[:, 0]
        beta[:, a, b, c] = sol[:, 1]
        
    int_field_000 = cumtrapz(tvec, np.dot(np.sqrt(
        kappa1), (alpha[:, 0, 0, 0] + beta[:, 0, 0, 0])), initial=0.0)
    int_field_011 = cumtrapz(tvec, np.dot(np.sqrt(
        kappa1), (alpha[:, 0, 1, 1] + beta[:, 0, 1, 1])), initial=0.0)

    figure_2b(kappa1,alpha,beta)
    sb.plt.show()